import cv2

cap = cv2.VideoCapture('../asset/cars.mp4')
bg_sub_mog = cv2.bgsegm.createBackgroundSubtractorMOG()
kernel = cv2.getStructuringElement(cv2.MORPH_RECT, (11, 11))
cars = []
car_n = 0

while cap.isOpened():
    ret, frame = cap.read()
    # H 1080, w 1920, 3
    # print(frame.shape)
    if ret:
        cv2.cvtColor(frame, cv2.COLOR_BGR2GRAY)
        blur = cv2.GaussianBlur(frame, (7, 7), 10)
        mask = bg_sub_mog.apply(blur)
        erode = cv2.erode(mask, kernel, 10)
        dilate = cv2.dilate(erode, kernel, 10)
        close = cv2.morphologyEx(dilate, cv2.MORPH_CLOSE, kernel)
        close = cv2.morphologyEx(close, cv2.MORPH_CLOSE, kernel, iterations=2)

        contours, h = cv2.findContours(close, cv2.RETR_TREE, cv2.CHAIN_APPROX_SIMPLE)
        cv2.line(frame, (10, 820), (1900, 820), (0, 255, 255), 3)
        for (i, c) in enumerate(contours):
            (x, y, w, h) = cv2.boundingRect(c)
            isShow = (w >= 90) and (h >= 90)
            if not isShow:
                continue

            cv2.rectangle(frame, (x, y), (x + w, y + h), (0, 0, 255), 2)
            center_p = (x + int(w / 2), y + int(h / 2))
            cars.append(center_p)
            print(center_p)
            cv2.circle(frame, center_p, 5, (0, 0, 255), -1)
            for (x, y) in cars:
                # 一辆车在y轴上的移动，每一帧图像检测，必须规定一个范围
                # 范围过大，有可能是同一辆车
                # 范围过小， 有的车统计不到
                if 864 < y < 878:
                    print("cars:"+str(y))
                    car_n += 1
                    cars.remove((x, y))
        cars.clear()
        cv2.putText(frame, "Cars Count:" + str(car_n), (500, 60), cv2.FONT_HERSHEY_SIMPLEX, 2, (0, 0, 255), 5)
        cv2.imshow('video', frame)
    key = cv2.waitKey(1)
    # esc 按钮
    if key & 0xFF == ord('q'):
        break

cap.release()
cv2.destroyAllWindows()
